from time import time
from typing import Callable
from pika.spec import Basic, BasicProperties
from pika.adapters.blocking_connection import BlockingChannel

from src.logger import Logger
from src.config import Config
from src.result_body import ResultBody


ConsumerCallback = Callable[
    [BlockingChannel, Basic.Deliver, BasicProperties, bytes], ResultBody
]


class Consumer:
    def __init__(
        self, queue: str, channel: BlockingChannel, callback: ConsumerCallback
    ) -> None:
        self.config = Config()

        channel.queue_declare(queue, durable=True)
        channel.basic_consume(queue, self._register_callback(callback), False)
        self.channel = channel
        self.result_queue = self.config.report.queue
        self.channel.queue_declare(self.result_queue, durable=True)

    BODY_TAG = "\x1b[1;36m body:\x1b[0m"

    def _register_callback(self, callback: ConsumerCallback):
        def _callback(
            channel: BlockingChannel,
            method: Basic.Deliver,
            props: BasicProperties,
            body: bytes,
        ) -> None:
            start = time()
            Logger.Info(
                f"message received, user: {props.user_id} app: {props.app_id}\n{self.BODY_TAG} {body}"
            )
            result = callback(channel, method, props, body)
            result.set_user_id(props.user_id)

            self.publish_result(result)
            channel.basic_ack(delivery_tag=method.delivery_tag)

            used_time = time() - start
            message = f"task finished, \x1b[1;33musing: +{used_time:.3f}s\x1b[0m\n{self.BODY_TAG} {result.to_json()}"
            # Todo 按 status 分级输出
            Logger.Info(message)
            return None

        return _callback

    def publish_result(self, body: ResultBody):
        self.channel.basic_publish(
            exchange="",
            routing_key=self.config.report.queue,
            body=body.to_bytes(),
            properties=BasicProperties(
                app_id=body.app_id,
                user_id=body.user_id,
                content_encoding="utf-8",
                content_type="application/json",
            ),
        )
